import pandas as pd
import requests
import time
import re
from urllib.parse import quote

def extract_cas_from_title(title):

    # Pattern for CAS number: numbers-numbers-number
    cas_pattern = r'\b\d+-\d+-\d+\b'
    match = re.search(cas_pattern, title)
    return match.group(0) if match else ''

def search_pubchem(compound_name):
    """
    Search PubChem for a compound and return its CID, CAS, and SMILES.
    Now with URL encoding for special characters.
    """
    try:
        # URL encode the compound name
        encoded_name = quote(compound_name)
        
        # First, search for the compound to get its CID
        search_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{encoded_name}/cids/JSON"
        response = requests.get(search_url)
        
        if response.status_code == 200:
            cid = str(response.json()['IdentifierList']['CID'][0])  # Get the best match
            
            # Get the title from the compound summary
            summary_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/synonyms/JSON"
            summary_response = requests.get(summary_url)
            
            cas = ''
            if summary_response.status_code == 200:
                synonyms = summary_response.json()['InformationList']['Information'][0]['Synonym']
                # Look for CAS number in synonyms
                for synonym in synonyms:
                    cas_match = extract_cas_from_title(synonym)
                    if cas_match:
                        cas = cas_match
                        break
            
            # Get SMILES using the CID
            property_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/property/CanonicalSMILES/JSON"
            prop_response = requests.get(property_url)
            
            if prop_response.status_code == 200:
                smiles = prop_response.json()['PropertyTable']['Properties'][0]['CanonicalSMILES']
                
                return {
                    'CID': cid,
                    'SMILES': smiles,
                    'CAS': cas
                }
        
        return {'CID': '', 'SMILES': '', 'CAS': ''}
    
    except Exception as e:
        print(f"Error processing {compound_name}: {str(e)}")
        return {'CID': '', 'SMILES': '', 'CAS': ''}

def update_excel_with_pubchem_data(input_file, output_file):

    # Read the Excel file
    df = pd.read_excel(input_file)
    df['CID'] = df['CID'].astype(str)
    df['CAS'] = df['CAS'].astype(str)
    df['SMILES'] = df['SMILES'].astype(str)
    
    # Get unique compound names
    unique_compounds = df['Name'].unique()
    print(f"Found {len(unique_compounds)} unique compounds out of {len(df)} total entries")
    
    # Create a dictionary to store results for each unique compound
    compound_data = {}
    
    # Process unique compounds
    total = len(unique_compounds)
    for idx, compound_name in enumerate(unique_compounds):
        compound_name = str(compound_name)
        print(f"Processing {idx+1}/{total}: {compound_name}")
        
        # Search PubChem
        result = search_pubchem(compound_name)
        compound_data[compound_name] = result
        
        # Print results for debugging
        print(f"Found data: CID={result['CID']}, CAS={result['CAS']}, SMILES={result['SMILES']}")
        
        # Add a small delay to avoid overwhelming the PubChem API
        time.sleep(0.2)
    
    # Fill in the existing columns using the cached results
    print("Filling in data for all entries...")
    for idx, row in df.iterrows():
        name = str(row['Name'])
        data = compound_data[name]
        df.at[idx, 'CID'] = data['CID']
        df.at[idx, 'CAS'] = data['CAS']
        df.at[idx, 'SMILES'] = data['SMILES']
    
    # Save the updated dataframe to the same file
    df.to_excel(output_file, index=False)
    print(f"Processing complete. Results saved to {output_file}")
    
    # Print some statistics
    successful_lookups = sum(1 for data in compound_data.values() if data['CID'])
    print(f"\nStatistics:")
    print(f"Total unique compounds processed: {total}")
    print(f"Successful lookups: {successful_lookups}")
    print(f"Failed lookups: {total - successful_lookups}")


# Usage
input_file = "Matched_DB.xlsx"  
output_file = "Final_API.xlsx"  

update_excel_with_pubchem_data(input_file, output_file)